package pos.trigram;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;

import com.thoughtworks.xstream.XStream;

public class Training {

	public static final String TAGSEPARATOR = ":";

	public static final String START = "START";

	// {tag, {word, count}}
	private HashMap<String, HashMap<String, Integer>> tagWordMap = new HashMap<String, HashMap<String, Integer>>();

	// word given tag probability.
	// {tag, {word, probability}} 
//	private HashMap<String, HashMap<String, Double>> tagWordProb = new HashMap<String, HashMap<String, Double>>();

	// keeping track of a tag given another tag (conditional probability).
	// {currentTag, {previousTag, count}
	private HashMap<String, HashMap<String, Integer>> tagTagMap = new HashMap<String, HashMap<String, Integer>>();

	// keeping track of a tag given another tag (conditional probability).
	// {currentTag, {previousTag+TAGSEPARATOR+secondPreviousTag, count}
	private HashMap<String, HashMap<String, Integer>> tagTag2Map = new HashMap<String, HashMap<String, Integer>>();

	// keeping track of a tag given another tag (conditional probability).
	// {currentTag, {previousTag, probability}
//	private HashMap<String, HashMap<String, Double>> tagTagProb = new HashMap<String, HashMap<String, Double>>();

	// {word, {tag, count}}
	private HashMap<String, HashMap<String, Integer>> wordTagMap = new HashMap<String, HashMap<String, Integer>>();

	// {word, MostFreqentTag}
	private HashMap<String, String> wordMap = new HashMap<String, String>();

	private HashMap<String, Integer> tagCount = new HashMap<String, Integer>();

	private HashMap<String, Integer> tag2Count = new HashMap<String, Integer>();

	private String mostUsedTag;

	public static void main (String[] args) {
		
		String filename = null;
		
		for(int i = 0; i < args.length; i = i+2) {
			if(args[i].equalsIgnoreCase("-f"))
				filename = args[i+1];
		}

		Training training = new Training();
		training.train(filename);
		HashMap<String, String> wordMap = training.getBaselineMap();
		System.out.println(wordMap);
		System.out.println(training.getMostUsedTag());
		training.serialize();
//		training.deserialize();
		//		System.out.println(training.getMostUsedTag());
//		System.out.println(training.getTagTagProb());
	}

	public void train(String filename){

		BufferedReader br = null;
		try {

			// reading the file and populating all the structures.
			{
				br = new BufferedReader(new FileReader(new File(filename)));
				String eachLine = null;
				String prevTag = START;
				String prev2Tag = START; // indicating t(i-2)
				tagCount.put(START, 1);
				while((eachLine=br.readLine())!=null){
					String[] list = eachLine.trim().split(" ");
					String tag = list [0];
					String word = list [1]/*.toLowerCase()*/;

					// To get the most used Tag.
					int count = 0;
					if(tagCount.containsKey(tag)) {
						count = tagCount.get(tag);
					}
					tagCount.put(tag, ++count);

					// TO keep counts of t(i-1) t(i-2) occurences.
					String prevTagCombined = combineTags(prevTag, prev2Tag); 
					count = 0;
					if(tag2Count.containsKey(prevTagCombined)) {
						count = tag2Count.get(prevTagCombined);
					}
					tag2Count.put(prevTagCombined, ++count);
					
					
					// To populate tagWordMap.
					if (null != tagWordMap && !tagWordMap.containsKey(tag)) {
						HashMap<String, Integer> wordmap = new HashMap<String, Integer>();
						tagWordMap.put(tag, wordmap);
					}

					if (tagWordMap.containsKey(tag)) {
						HashMap<String, Integer> wordmap = tagWordMap.get(tag);
						if (null != wordmap && !wordmap.isEmpty()
								&& wordmap.containsKey(word)) {
							int wordCount = wordmap.get(word);
							wordmap.put(word, wordCount + 1);
						} else {
							wordmap.put(word, 1);
						}
					} else {
						System.err.println("Some problem");
					}

					// To populate wordTagMap.
					HashMap<String, Integer> tagMap;
					if(wordTagMap.containsKey(word)) {
						tagMap = wordTagMap.get(word);
					} 
					else {
						tagMap = new HashMap<String, Integer>();
					}
					count = 0;
					if(tagMap.containsKey(tag)) {
						count = tagMap.get(tag);
					}

					tagMap.put(tag, ++count);
					wordTagMap.put(word, tagMap);

					//
					HashMap<String, Integer> prevTagMap = tagTagMap.get(tag);
					if(prevTagMap == null) {
						prevTagMap = new HashMap<String, Integer>();
					}
					count = 0;
					if(prevTagMap.containsKey(prevTag)) {
						count = prevTagMap.get(prevTag);
					}
					prevTagMap.put(prevTag, ++count);
					tagTagMap.put(tag, prevTagMap);

					prevTagMap = tagTag2Map.get(tag);
					if(prevTagMap == null) {
						prevTagMap = new HashMap<String, Integer>();
					}
					count = 0;
					if(prevTagMap.containsKey(prevTagCombined)) {
						count = prevTagMap.get(prevTagCombined);
					}
					prevTagMap.put(prevTagCombined, ++count);
					tagTag2Map.put(tag, prevTagMap);
					prev2Tag = prevTag;
					prevTag = tag;
				}
			}

//			normalize();

//			tagTagProb = convertTagTagProb(tagTagMap, tagCount);

//			tagWordProb = convertTagWordProb(tagWordMap, tagCount);

			mostUsedTag = findMostUsedTag();

			wordMap = findMostFrequentTag();

		} catch (Exception e) {
			e.printStackTrace();
		} finally{
			if(br!=null){
				try {
					br.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
	}

	public static String combineTags(String prevTag, String prev2Tag) {
		return prevTag + TAGSEPARATOR + prev2Tag;
	}

	/**
	 * finds the most used tag for each word.
	 * @return 
	 */
	private HashMap<String, String> findMostFrequentTag() {
		Set<String> words = wordTagMap.keySet();
		HashMap<String, String> wordMap = new HashMap<String, String>();
		for (String word : words) {
			HashMap<String, Integer> tagMap = wordTagMap.get(word);
			int max = 0;
			String bestTag = null;		// signifies the most frequent tag for a given word
			// finding the most freqent tag.
			Set<String> tagSet = tagMap.keySet();
			for (String tag : tagSet) {
				int count = tagMap.get(tag);
				if(count > max) {
					max = count;
					bestTag = tag;
				}
			}

			wordMap.put(word, bestTag);
		}
		return wordMap;
	}

	private String findMostUsedTag() {
		// Finding the most used Tag.
		int mostUsedTagCount = 0;
		Set<String> tags = tagCount.keySet();
		String mostUsedTag = null;
		for (String tag : tags) {
			if(tagCount.get(tag) > mostUsedTagCount) {
				mostUsedTagCount = tagCount.get(tag);
				mostUsedTag = tag;
			}
		}
		return mostUsedTag;
	}

	private HashMap<String,HashMap<String,Double>> convertTagWordProb(HashMap<String, HashMap<String,Integer>> tagWordMap, HashMap<String, Integer> tagCount) {
		// converting counts for word|tag to probability of word|tag
		HashMap<String, HashMap<String,Double>> tagWordProb = new HashMap<String, HashMap<String,Double>>();
		Set<String> tagSet = tagWordMap.keySet();
		for (String tag : tagSet) {
			HashMap<String, Integer> wordMap = tagWordMap.get(tag);
			Integer tagTotalCount = tagCount.get(tag);
			Set<String> words = wordMap.keySet();
			HashMap<String, Double> wordProbMap = new HashMap<String, Double>();
			for (String word : words) {
				wordProbMap.put(word, wordMap.get(word) / (double)tagTotalCount);
			}
			tagWordProb.put(tag, wordProbMap);
		}
		return tagWordProb;
	}

	private HashMap<String,HashMap<String,Double>> convertTagTagProb(HashMap<String, HashMap<String, Integer>> tagTagMap2, HashMap<String, Integer> tagCount2) {
		// converting counts for tag|tag to probability of tag|tag
		Set<String> tagSet = tagTag2Map.keySet();
		HashMap<String,HashMap<String,Double>> tagTagProb = new HashMap<String, HashMap<String,Double>>();
		for (String tag : tagSet) {
			HashMap<String, Integer> prevTagMap = tagTag2Map.get(tag);
			Set<String> prevTags = prevTagMap.keySet();
			HashMap<String, Double> prevTagProbMap = new HashMap<String, Double>();
			for (String prevTag : prevTags) {
				prevTagProbMap.put(prevTag, prevTagMap.get(prevTag) / (double)tagCount.get(prevTag));
			}
			tagTagProb.put(tag, prevTagProbMap);
		}
		return tagTagProb;
	}

	public HashMap<String, HashMap<String, Integer>> getTagWordMap() {
		return tagWordMap;
	}

	public HashMap<String, HashMap<String, Integer>> getWordTagMap() {
		return wordTagMap;
	}

	public HashMap<String, String> getBaselineMap() {
		return wordMap;
	}

	public String getMostUsedTag() {
		return mostUsedTag;
	}

	public HashMap<String,HashMap<String,Integer>> getTagTagMap() {
		return tagTag2Map;
	}

	public HashMap<String,Integer> getTagCount() {
		return tagCount;
	}

	public HashMap<String,Integer> getTag2Count() {
		return tag2Count;
	}

	public Set<String> getWordSet() {
		return wordMap.keySet();
	}

	public void serialize() {
		try {
			XStream xstream = new XStream();

			serialize(xstream, "tagWordMap", tagWordMap);
			serialize(xstream, "tagTagMap", tagTag2Map);
//			serialize(xstream, "tagTagProb", tagTagProb);
//			serialize(xstream, "tagWordProb", tagWordProb);
			serialize(xstream, "mostUsedTag", mostUsedTag);
			serialize(xstream, "tagCount", tagCount);
			serialize(xstream, "tag2Count", tag2Count);
			serialize(xstream, "wordTagMap", wordTagMap);
			serialize(xstream, "wordMap", wordMap);
		} catch (IOException e) {
			e.printStackTrace();
		}


	}

	public void deserialize() {
		XStream xstream = new XStream();
		tagTag2Map = (HashMap<String, HashMap<String, Integer>>)deserialize(xstream, "tagTagMap");
		tagWordMap = (HashMap<String, HashMap<String, Integer>>)deserialize(xstream, "tagWordMap");
//		tagTagProb = (HashMap<String, HashMap<String,Double>>)deserialize(xstream, "tagTagProb");
//		tagWordProb = (HashMap<String, HashMap<String,Double>>)deserialize(xstream, "tagWordProb");
		mostUsedTag = (String)deserialize(xstream, "mostUsedTag");
		tagCount = (HashMap<String, Integer>)deserialize(xstream, "tagCount");
		tag2Count = (HashMap<String, Integer>)deserialize(xstream, "tag2Count");
		wordTagMap = (HashMap<String, HashMap<String, Integer>>)deserialize(xstream, "wordTagMap");
		wordMap = (HashMap<String, String>) deserialize(xstream, "wordMap");
	}

	private void serialize(XStream xstream, String filename, Object obj) throws IOException {
		BufferedWriter writer = new BufferedWriter(new FileWriter(filename));
		writer.write(xstream.toXML(obj));
		writer.flush();
		writer.close();
	}

	public void normalize() {

		// add 1 smoothing to Tag|previousTag.
		Set<String> tagSet = tagCount.keySet();
		for (String tag : tagSet) {
			HashSet<String> tags = copySet(tagSet);;
			HashMap<String, Integer> hashMap = tagTag2Map.get(tag);
			if(hashMap != null) {
				Set<String> keySet = hashMap.keySet();
				tags.removeAll(keySet);
				for (String key : keySet) {
					hashMap.put(key, hashMap.get(key)+1);
				}
			}
			else
				hashMap = new HashMap<String, Integer>();

			for (String string : tags) {
				hashMap.put(string, 1);
			}
			tagTag2Map.put(tag, hashMap);
		}


		// add 1 smoothing to word|Tag.
		Set<String> wordSet = wordTagMap.keySet();
		for (String tag : tagSet) {
			HashSet<String> words = copySet(wordSet);
			HashMap<String, Integer> hashMap = tagWordMap.get(tag);
			if(hashMap!=null) {
				words.removeAll(hashMap.keySet());
				for (String key : hashMap.keySet()) {
					hashMap.put(key, hashMap.get(key)+1);
				}
			}
			else 
				hashMap = new HashMap<String, Integer>();
			
			for (String string : words) {
				hashMap.put(string, 1);
			}
			tagWordMap.put(tag, hashMap);
			tagCount.put(tag, tagCount.get(tag) + wordSet.size());
		}
	}

	private HashSet<String> copySet(Set<String> tagSet) {
		HashSet<String> set = new HashSet<String>();
		for (String string : tagSet) {
			set.add(string);
		}
		return set;
	}

	private Object deserialize(XStream xstream, String filename) {
		try {
			BufferedReader reader = new BufferedReader(new FileReader(filename));
			StringBuffer buffer = new StringBuffer();
			String line;
			while((line = reader.readLine()) != null) {
				buffer.append(line);
			}
			return xstream.fromXML(buffer.toString());
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		return null;
	}
}
